# 截屏切分模块
# 为后续的分析提供更专一的数据源
import cv2
import time
import math
import numpy as np
import logging
from shape_judgment.screen_region_storage import storage_a_region,storage_a_region_pos_only,get_region_pos_only

def split_by_vertical_line(screen_shot_img):
    # return 返回切完之后的图像列表

    # time1 = time.time()

    # 1. 开始分析竖线:
    split_pos_list = [] # 记录用于分割的x坐标
    height, width = screen_shot_img.shape[:2]
    for x in range(width):
        # if x<4 :continue 最边缘也要判断,用于把区域截出来, 不然之后截出来的区域会有这(60,60,60)的边框
        # 一列的开始:
        number_matchs = 0 # 统计有多少个像素是匹配的
        for y in range(height):
            if (screen_shot_img[y, x][0] == 60) and (screen_shot_img[y, x][1] == 60) and (screen_shot_img[y, x][2] == 60):
                number_matchs += 1
        # 一列的统计的结束:
        if number_matchs > height/2: # 竖线一般至少10个像素
            split_pos_list.append(x)

    # print("# 65 分析出用于分割的竖线:", split_pos_list)
    # time2 = time.time()
    # print(f"找竖线用时:{time2-time1}")
    # # DEBUG 保存, 用于调试
    # cv2.imwrite("out_test_screen_shot.png", screen_shot_img)

    # 2. 分割, 并存到列表中
    # 刚好没截取头尾, 正是需求上需要的, 不过逻辑不严密
    images = []
    for i in range(len(split_pos_list)-1):
        if split_pos_list[i+1]-(split_pos_list[i]+1) <=25 :continue
        img = screen_shot_img[:, split_pos_list[i]+1:split_pos_list[i+1]]
        # cv2.imwrite(f"out_test_screen_shot_{split_pos_list[i]+1}_{split_pos_list[i+1]}.png", img)
        images.append(([split_pos_list[i]+1,split_pos_list[i+1]],img))
    
    return images

# 在图表区内, 取更暗的颜色(35,35,35), 这是一些垂直的线, 把图像分割成几块
def split_by_vertical_line_35(screen_shot_img):

    # 1. 开始分析竖线:
    split_pos_list = [0] # 记录用于分割的x坐标
    height, width = screen_shot_img.shape[:2]
    for x in range(width):
        # 一列的开始:
        number_matchs = 0 # 统计有多少个像素是匹配的
        for y in range(height):
            if (screen_shot_img[y, x][0] == 35) and (screen_shot_img[y, x][1] == 35) and (screen_shot_img[y, x][2] == 35):
                number_matchs += 1
        # 一列的统计的结束:
        if number_matchs > height/2: # 竖线一般至少10个像素
            split_pos_list.append(x)

    split_pos_list.append(width) # 保证头尾也切进去
    # print("# 97 第二次竖线分割结果:",split_pos_list)

    # 2. 分割, 并存到列表中
    images = []

    for i in range(len(split_pos_list)-1):
        if split_pos_list[i+1]-(split_pos_list[i]+1) <=0 :continue
        img = screen_shot_img[:, split_pos_list[i]+1:split_pos_list[i+1]]
        # cv2.imwrite(f"out_test_screen_shot_[inner_chart]_{split_pos_list[i]+1}_{split_pos_list[i+1]}.png", img)
        images.append(([split_pos_list[i]+1,split_pos_list[i+1]] ,img))

    return images

# 根据横线分割图片 用于筹码
def split_by_horizontal_line(screen_shot_img):
    # 横线将同时有多种颜色, 因为之前单单用竖线分割, 没有割掉菜单. 
    # return 列表, 列表中的元素是分割后的图像

    # 1. 找到分割点:
    split_pos_list = [] # 存储分割点
    height, width = screen_shot_img.shape[:2]
    for y in range(height):
        # 一行的开始:
        number_matchs = 0 # 统计有多少个像素是匹配的
        for x in range(width):
            if (screen_shot_img[y, x][0] == 35) and (screen_shot_img[y, x][1] == 35) and (screen_shot_img[y, x][2] == 35)\
                or (screen_shot_img[y, x][0] == 60) and (screen_shot_img[y, x][1] == 60) and (screen_shot_img[y, x][2] == 60)\
                or (screen_shot_img[y, x][0] == 74) and (screen_shot_img[y, x][1] == 74) and (screen_shot_img[y, x][2] == 74)\
                :
                number_matchs += 1
        if number_matchs > width/2: 
            split_pos_list.append(y)
    
    # print("# 131 用于分割的y坐标:",split_pos_list)
    # 这次不管头尾了, 因为头尾都有UI, 只想取图表

    images = []
    biggest_img = None
    biggest_img_size = 0
    for i in range(len(split_pos_list)-1):
        if split_pos_list[i+1]-(split_pos_list[i]+1) <=0 :continue
        if biggest_img_size < split_pos_list[i+1]-(split_pos_list[i]+1) :
            biggest_img = ([0,width,split_pos_list[i]+1,split_pos_list[i+1]], screen_shot_img[split_pos_list[i]+1:split_pos_list[i+1], :]         )
            biggest_img_size = split_pos_list[i+1]-(split_pos_list[i]+1)
        # cv2.imwrite(f"out_test_screen_shot_[inner_chart_2]_{split_pos_list[i]+1}_{split_pos_list[i+1]}.png", img)
    
    images.append(biggest_img)
    return images



# 根据横线分割图片 用于大图表
def split_by_horizontal_line_35_38_60(screen_shot_img):
    # 横线将同时有多种颜色, 因为之前单单用竖线分割, 没有割掉菜单. 
    # return 列表, 列表中的元素是分割后的图像

    # 1. 找到分割点:
    split_pos_list = [] # 存储分割点
    height, width = screen_shot_img.shape[:2]
    for y in range(height):
        # 一行的开始:
        number_matchs = 0 # 统计有多少个像素是匹配的
        for x in range(width):
            if (screen_shot_img[y, x][0] == 35) and (screen_shot_img[y, x][1] == 35) and (screen_shot_img[y, x][2] == 35)\
                or (screen_shot_img[y, x][0] == 38) and (screen_shot_img[y, x][1] == 38) and (screen_shot_img[y, x][2] == 38)\
                or (screen_shot_img[y, x][0] == 60) and (screen_shot_img[y, x][1] == 60) and (screen_shot_img[y, x][2] == 60)\
                or (screen_shot_img[y, x][0] == 74) and (screen_shot_img[y, x][1] == 74) and (screen_shot_img[y, x][2] == 74)\
                :
                number_matchs += 1
        if number_matchs > width/2: 
            split_pos_list.append(y)
    
    # print("# 131 用于分割的y坐标:",split_pos_list)
    # 这次不管头尾了, 因为头尾都有UI, 只想取图表

    images = []
    # biggest_img = None
    # biggest_img_size = 0
    for i in range(len(split_pos_list)-1):
        if split_pos_list[i+1]-(split_pos_list[i]+1) <=0 :continue
        # if biggest_img_size < split_pos_list[i+1]-(split_pos_list[i]+1) :
        images.append(([0,width,split_pos_list[i]+1,split_pos_list[i+1]]
                      , screen_shot_img[split_pos_list[i]+1:split_pos_list[i+1], :]))
        # biggest_img_size = split_pos_list[i+1]-(split_pos_list[i]+1)
        # cv2.imwrite(f"out_test_screen_shot_[inner_chart_2]_{split_pos_list[i]+1}_{split_pos_list[i+1]}.png", img)
    
    
    return images





# 可能会有多个匹配, 返回所有匹配先
def find_all_matches(screen_img, template_img, threshold):
    """获取所有符合阈值的匹配结果，返回列表[(x, y, w, h, 匹配度)]"""

    h, w = template_img.shape[:2]
    
    if h == 0 or w == 0:
        return []
    
    if screen_img.size < template_img.size:  # 截图尺寸小于模板尺寸
        print("截图尺寸小于模板尺寸")
        return []
    
    # 执行模板匹配并获取匹配度矩阵
    result = cv2.matchTemplate(
        screen_img, 
        template_img, 
        cv2.TM_CCOEFF_NORMED
    )
    # 找到所有超过阈值的位置
    locations = np.where(result >= threshold)
    matches = []
    
    # 遍历所有匹配位置，记录坐标和匹配度
    for y, x in zip(locations[0], locations[1]):
        match_value = result[y, x]  # 该位置的匹配度（0-1）
        matches.append((x, y, w, h, match_value))  # 存储x坐标、y坐标、宽、高、匹配度
    return matches

# 找 模板匹配
def find_template_position(screen_img, template_img, threshold=0.85):
    """
    在屏幕截图中查找模板位置
    :param screen_img: 全屏截图（BGR格式）
    :param template_img: 模板图片（BGR格式）
    :param threshold: 匹配阈值（0-1，越高越严格） 数值为0.8时,还是会有错误, 比如需匹配0.00%,匹配到的却是0.05%
    :return: 匹配区域的左上角坐标(x, y)，若未找到返回None
    """
    
    all_matches  = find_all_matches(screen_img, template_img,threshold)
    if not all_matches:
        return None  # 无匹配结果
    # 按X坐标升序排序（X越小越靠左），取第一个
    # 若X相同，可选匹配度最高的
    all_matches.sort(key=lambda m: (-m[4]))  # 优先按X排序，再按y排序 找到左上角的一个叉按钮
    rightmost = all_matches[0]
    return (rightmost[0], rightmost[1], rightmost[2], rightmost[3])  # 返回(x, y, w, h)

# 获取关键的图像
def get_key_image(screen_shot_img):

    images_1 = split_by_vertical_line(screen_shot_img)


    # BEGIN 防范
    if 0 == len(images_1): # 竖线分割失败
        return []
    # END 防范
    
    chart_area = images_1[0][1] # 用于取最大的图
    obj_image_biggest=images_1[0]
    for i in range(len(images_1)):
        if images_1[i][1].shape[1] > chart_area.shape[1]: # 宽度最大的图
            chart_area = images_1[i][1]
            obj_image_biggest = images_1[i]
    storage_a_region(obj_image_biggest[0][0],obj_image_biggest[0][1],0,obj_image_biggest[1].shape[0],"最大图表区")

    # 遍历第二次, 找出两个侧边栏
    int_number_right = 0
    obj_img_right_side_1 = None
    obj_img_right_side_2 = None # 右侧边栏2
    for i in range(len(images_1)):
        # 当在图表的右侧:
        if images_1[i][0][0] > obj_image_biggest[0][0]:
            # 判断是分隔线:
            if(images_1[i][1].shape[1] < 28):continue
            int_number_right+=1
            # print(f"# 172 右侧边栏{int_number_right}:  {images_1[i][0][0]},{images_1[i][0][1]} \t {0},{images_1[i][1].shape[0]}")
            storage_a_region(images_1[i][0][0],images_1[i][0][1],0,images_1[i][1].shape[0],f"右侧边栏{int_number_right}")
            if(1 == int_number_right): obj_img_right_side_1=images_1[i]
            if(2 == int_number_right): obj_img_right_side_2=images_1[i]

    # 刚返回的是相对位置
    images_vertical_2 = split_by_vertical_line_35(chart_area)
    for i in range(len(images_vertical_2)):
        # print("# 179", images_vertical_2[i][0])
        images_vertical_2[i][0][0]+=obj_image_biggest[0][0]
        images_vertical_2[i][0][1]+=obj_image_biggest[0][0]
    
    if images_vertical_2 is None or 1 >= len(images_vertical_2): # 竖线分割失败
        return []

    # 取筹码图:
    image_chip = images_vertical_2[2]

    # 切出筹码图:
    image_clips_for_chip = split_by_horizontal_line(image_chip[1])

    # 从相对坐标处理成全局坐标:
    for i in range(len(image_clips_for_chip)):
        image_clips_for_chip[i][0][0]+=image_chip[0][0]
        image_clips_for_chip[i][0][1]+=image_chip[0][0]
        storage_a_region(image_clips_for_chip[i][0][0]
                         ,image_clips_for_chip[i][0][1]
                         ,image_clips_for_chip[i][0][2]
                         ,image_clips_for_chip[i][0][3]
                         ,"筹码区")
    
    # 切分K线图区域 切出顶部菜单栏, 用于判断有没有筹码这个按钮
    # 切分出顶部三个菜单与底部4个菜单栏
    # 这里不能用 images_vertical_2[0][1] 因为裁掉筹码区后, 筹码按钮就会受到切割. 
    images_main_chart_clips = split_by_horizontal_line_35_38_60(obj_image_biggest[1])
    for i in range(len(images_main_chart_clips)):
        if(0==i):continue
        if(images_main_chart_clips[i][1].shape[0] < 14):continue
        # cv2.imwrite(f"out_put_test_screen_shot_split_[main_chart]_{i}.png", images_main_chart_clips[i][1])
        arr_pos_info = images_main_chart_clips[i][0]
        arr_pos_info[0]+=obj_image_biggest[0][0]
        arr_pos_info[1]+=obj_image_biggest[0][0]
        storage_a_region(arr_pos_info[0],arr_pos_info[1],arr_pos_info[2],arr_pos_info[3],"筹码按钮所在的菜单栏")
        break

    # BEGIN 切出筹码区之外的图表区:
    region_chat_area = get_region_pos_only("最大图表区")
    region_chat_chip = get_region_pos_only("筹码区")
    x1 = region_chat_area[0]
    x2 = region_chat_chip[0]
    y1 = region_chat_chip[2]
    y2 = region_chat_chip[3]
    # cv2.imwrite("screen_shot_284.png", screen_shot_img[y1:y2,x1:x2]) # 正确 
    images_split_k = split_by_horizontal_line_35_38_60(screen_shot_img[y1:y2,x1:x2])
    image_obj_k_larger = images_split_k[0]
    image_obj_k_larger_height = image_obj_k_larger[0][3]-image_obj_k_larger[0][2]
    for i in range(len(images_split_k)):
        # cv2.imwrite(f"screen_shot_287_{i}.png", images_split_k[i][1]) # 输出5张图, 其中,第零个是按钮栏
        arr_pos_info = images_split_k[i][0]
        img_height_inner = arr_pos_info[3]-arr_pos_info[2]
        if(image_obj_k_larger_height < img_height_inner):
            image_obj_k_larger_height = img_height_inner
            image_obj_k_larger = images_split_k[i]
    # cv2.imwrite(f"screen_shot_k_larger.png", image_obj_k_larger[1]) # 这一步正确 但右侧有一排看不懂的数字
    # 切出K线图表区,用最大的图去切, 把右边那一列没用的去掉
    images_k_larger = split_by_vertical_line_35(image_obj_k_larger[1])
    img_k_larger_2 = None
    # for i in range(len(images_k_larger)):
    #     cv2.imwrite(f"screen_shot_k_larger_split_{i}.png",images_k_larger[i][1])
    #     print("# 300")
    # print(f"# 301 images size:{len(images_k_larger)}")
    if 2 == len(images_k_larger):
        img_k_larger_2 = images_k_larger[0]
    else:
        raise Exception("!!!!    ERR 306 K线图表区切分失败")
    obj_img_larger_k = ([x1,x1+img_k_larger_2[0][1],y1+image_obj_k_larger[0][2],y1+image_obj_k_larger[0][3]]
                        ,screen_shot_img[y1+image_obj_k_larger[0][2]:y1+image_obj_k_larger[0][3]
                                         ,x1:x1+img_k_larger_2[0][1]]) # 再裁剪一次,确保坐标算式正确
    cv2.imwrite(f"screen_shot_k_larger_split_310.png",obj_img_larger_k[1]) # 正确
    storage_a_region(obj_img_larger_k[0][0],obj_img_larger_k[0][1],obj_img_larger_k[0][2],obj_img_larger_k[0][3],"K线图表区")




    # END 切出筹码区之外的图表区
    if (obj_img_right_side_2 is not None):
        image_list_right_side_2 = split_by_horizontal_line_35_38_60(obj_img_right_side_2[1])
        # print(" # 277 image_list_right_side_2:",image_list_right_side_2)
        for i in range(len(image_list_right_side_2)):
            if(image_list_right_side_2[i][1].shape[0] <= 22):continue
            print("# 228 height:",image_list_right_side_2[i][1].shape[0])
            arr_pos_info = image_list_right_side_2[i][0]
            arr_pos_info[0]+=obj_img_right_side_2[0][0]
            arr_pos_info[1]+=obj_img_right_side_2[0][0]
            storage_a_region(arr_pos_info[0],arr_pos_info[1],arr_pos_info[2],arr_pos_info[3],"右侧边栏2的顶部按钮区")
            # cv2.imwrite(f"out_put_test_screen_shot_split_[right_side_2]_{i}.png", image_list_right_side_2[i][1])
            break
    else:
        cv2.imwrite(f"ERR_331_screen_shot_right_side_2_not_found.png", screen_shot_img)        
        cv2.imwrite(f"ERR_331_screen_shot_right_side_2_not_found_obj_image_biggest.png", obj_image_biggest[1])
        if obj_img_right_side_1:cv2.imwrite(f"ERR_331_screen_shot_right_side_2_not_found_obj_img_right_side_1.png", obj_img_right_side_1[1])
        for i in range(len(images_1)):
            cv2.imwrite(f"ERR_331_screen_shot_right_side_2_not_found_obj_img_right_side_1_[images_1]_{i}.png", images_1[i][1])
        raise Exception(f"!!!!    ERR 331 右侧边栏2切分,可能是窗口伸到右侧屏幕外  图像保存为 ERR_331_screen_shot_right_side_2_not_found.png")
    
    # BEGIN 右侧边栏1的处理
    # 切分出交易状态的栏位, 没有交易状态的股票, 就不处理
    if(obj_img_right_side_1 is not None):
        images_trading_status_clips = split_by_horizontal_line_35_38_60(obj_img_right_side_1[1])
        number_of_target = 3 # 第一个,是菜单栏 第二个, 是股票ID以及股票名称, 第三个就是目标[交易状态]
        number_of_target_2 = 2 # 是股票ID以及股票名称
        number_count = 0
        image_stock_name_and_price = None
        for i in range(len(images_trading_status_clips)):
            # cv2.imwrite(f"out_put_test_screen_shot_split_[right_side_1]_{i}.png", images_trading_status_clips[i][1])
            if 0 == i: continue
            if images_trading_status_clips[i][1].shape[0] < 14:continue
            number_count += 1
            if(number_count==number_of_target):
                # cv2.imwrite(f"out_put_test_screen_shot_split_[right_side_1]_{i}.png", images_trading_status_clips[i][1])
                x1,x2,y1,y2 = images_trading_status_clips[i][0]
                x1 = x1+obj_img_right_side_1[0][0]
                x2 = x2+obj_img_right_side_1[0][0]
                storage_a_region(x1,x2,y1,y2,"交易状态一整栏")# 用于判断有没有交易状态这一栏
                x1 += math.ceil((x2-x1)*.5)
                storage_a_region(x1,x2,y1,y2,"交易状态后半栏")# 用于判断是否退市
                break
            elif(number_of_target_2 == number_count):
                x1,x2,y1,y2 = images_trading_status_clips[i][0]
                x1 = x1+obj_img_right_side_1[0][0]
                x2 = x2+obj_img_right_side_1[0][0]
                storage_a_region(x1,x2,y1,y2,"股票名称与价格")
                image_stock_name_and_price = images_trading_status_clips[i]
        # 取右箭头按钮坐标:
        if image_stock_name_and_price is not None:
            template_img_right = cv2.imread("eastmoney desktop software picture/template_triangle_button_right.png")
            pos = find_template_position(image_stock_name_and_price[1], template_img_right)
            x,y,w,h = pos
            storage_a_region_pos_only(x,x+w, y,y+h, "右箭头按钮1") # 这里只记录一个右箭头按钮可能的位置, 因为同一时间, 页面的样式只有一种
            logging.info(f"get_key_image # 364 右箭头坐标:{pos}")
            print("# 321 BEGIN Debug 进一步截取,还能否匹配?")
            img_debug = image_stock_name_and_price[1][y:y+h,x:x+w]
            pos_debug = find_template_position(img_debug, template_img_right)
            if(pos_debug is not None):
                print("# 325 Debug 进一步截取,能匹配!")
            else:
                print("# 325 Debug 进一步截取,不能匹配!")
            print("# 321 END Debug 进一步截取,还能否匹配?")
    # END 右侧边栏1的处理


    return [image_clips_for_chip[0][1]]